package hw2.nsl;

import hw2.sign.FileIO;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

public class Bob {
	
	static PublicKey alicePub;
	static PrivateKey bobPriv;
	static final int portno = 7002;
	static final int malportno = 7001;
	static boolean encryptMode = true;
	static Integer myNonce = 0;

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		System.out.println("\nYou are running the Bob program.");
		if (args.length > 0){
			if (args[0].equalsIgnoreCase("plain")){
				encryptMode = false;
				System.out.println("Encryption has been turned off.");
			}
		}
		readKeys();
		ServerSocket sock = null;
		try {
			sock = new ServerSocket(Bob.portno);
		} catch (IOException e1) {
			e1.printStackTrace();
			System.exit(-1);
		}
		while( true ) {
			try{
				try{
					Socket client = sock.accept();
					ObjectInputStream in = new ObjectInputStream(client.getInputStream());

					NetworkMessage msg = (NetworkMessage)in.readObject();
					System.out.println("Just received: " + msg.toHexString());
					client.close();

					Bob.handleMessage(msg);
				}catch(Exception e){
					e.printStackTrace();
				}				
        	}catch(Exception e){
        		e.printStackTrace();
        	}
        }
	}
	
	public static void sendMessage(NetworkMessage m) throws Exception{
		Socket clientSocket = new Socket("127.0.0.1",malportno);
//		Socket clientSocket = new Socket("127.0.0.1",7000);
		ObjectOutputStream out = new ObjectOutputStream(clientSocket.getOutputStream());
		System.out.println("About to send: " + m.toHexString());
        out.writeObject(m);
		
		out.flush();
	}
	
	/**
	 * Read in Bob's public key and alice's private key
	 */
	private static void readKeys() {
		try {
			byte[] b = FileIO.getBytesFromFile("bob_priv.bin");
			PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(b);
			bobPriv =  KeyFactory.getInstance( "RSA" ).generatePrivate( spec );

			b = FileIO.getBytesFromFile("alice_pub.bin");
			X509EncodedKeySpec spec2 = new X509EncodedKeySpec(b);
			alicePub =  KeyFactory.getInstance( "RSA" ).generatePublic( spec2 );
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(-1);
		}
	}
	
	/**
	 * Enc with bob's public key and get the enc bytes back
	 * @param message
	 * @return
	 * @throws NoSuchAlgorithmException
	 * @throws NoSuchPaddingException
	 * @throws InvalidKeyException
	 * @throws IllegalBlockSizeException
	 * @throws BadPaddingException
	 */
	private static byte[] encryptToAlice(byte[] message) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
		Cipher enc = Cipher.getInstance("RSA/ECB/OAEPWITHSHA-256ANDMGF1PADDING");
		enc.init(Cipher.ENCRYPT_MODE, alicePub);
		return enc.doFinal(message);		
	}
	
	/**
	 * Decrypt our message
	 * @param message - bytes which were enced with our (Alice's) public key
	 * @return
	 * @throws NoSuchPaddingException 
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeyException 
	 * @throws BadPaddingException 
	 * @throws IllegalBlockSizeException 
	 */
	private static byte[] dec(byte[] message) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
		Cipher dec = Cipher.getInstance("RSA/ECB/OAEPWITHSHA-256ANDMGF1PADDING");
		dec.init(Cipher.DECRYPT_MODE, bobPriv);
		return dec.doFinal(message);
	}

	public static void handleMessage(NetworkMessage msg) throws Exception {
		byte[] text = msg.text;
		if (encryptMode){
			try{
				text = dec(text);
			}catch(Exception up){
				System.err.println("\nERROR: Failed to decrypt message");
				return;
			}
		}
		NetworkMessage m = new NetworkMessage(text, "");
		String[] vals = m.parseFields();
		switch(vals.length){
		case 1:
			// Bob believes it to be a message in part 3 of protocol
			// Bob checks against nonce and protocol ends
			if (myNonce == Integer.parseInt(vals[0])){
				System.out.println("\nSUCESS: I accept that I am talking to Alice!");
			}
			else{
				System.err.println("ERROR: Nonce didn't match! Rejecting connection");
			}
			break;
		case 2:
			// Bob thinks it is a message from part 1 of the protocol
			// Bob generates a nonce and sends it to alice with her nonce
			// vals[0] => Alice
			// vals[1] => N_a
			
			// create a nonce for myself
			SecureRandom sr = SecureRandom.getInstance("SHA1PRNG");
			myNonce = sr.nextInt();
			
			//send {Bob, N_a, N_b}K_a to alice
			String strmsg = "Bob " + vals[1] + " " + myNonce;
			byte[] bytes = strmsg.getBytes();
			if (encryptMode){
				bytes = encryptToAlice(bytes); 
			}
			NetworkMessage sendme = new NetworkMessage(bytes, "Bob");
			
			sendMessage(sendme);
			break;
		default:
			System.err.println("\nERROR: Received message with " + vals.length + " fields.");
		}
	}
}

